import os
import sys
import time
import glob
from PIL import Image
import numpy as np
import pydicom
import functools
from pydicom.tag import Tag
# from scipy.misc import imresize, imread, imshow
# import SimpleITK as sitk
import numpy as sitk

__all__ = [# 'cv_show',
           'read_dcm_by_pydicom',
            'read_dcm_info',
           'read_dcm_to_norm',
           'convert_pixel_to_ct',
           'read_dcm_to_8bit',
           'convert_dcm_grayscale8',
           'load_mhd_image',
           'normalize_ct_lung_to_unit',
           'get_numpy_from_sitk',
           'read_dcm_to_sitk']

PHOTOMETRIC_INTERPRETATION_TAG = (0x0028,0x0004)
BITS_STORED_TAG = (0x0028, 0x0101)
RESCALE_SLOPE_TAG = (0x0028,0x1053)
RESCALE_INTERCEPT_TAG = (0x0028,0x1052)
BODY_PART_TAG = (0x0018, 0x0015)
WINDOW_CENTER = (0x0028, 0x1050)
WINDOW_WIDTH = (0x0028, 0x1051)
PIXEL_SPACING_TAG = (0x0028,0x0030)


BODY_PART_NAME = ('BREAST', 'LSPINE', 'LSPINE', 'SKULL', 'ANKLE', 'HAND', 'KNEE',
                  'ABDOMEN', 'CSPINE', 'C-SPINE','SHOULDER', 'FOOT', 'CHEST',
                  'HIP', 'THORAX', 'COCCYX', 'ELBOW', 'LEG')

CHEST_TAG_NAME = ('胸', 'CHEST', 'THORAX')

sys_is_little_endian = (sys.byteorder == 'little')


def scipy_resize(image, size=(64, 64)):
    # return imresize(image, size)
    # resize_scale = (size[0]/image.shape[0], size[1]/image.shape[1])
    return np.array(Image.fromarray(image).resize(size))  # , resize_scale
# def cv_show(name, image, resize=1):
#     H,W = image.shape[0:2]
#     cv2.namedWindow(name, cv2.WINDOW_NORMAL)
#     cv2.imshow(name, image.astype(np.uint8))
#     cv2.resizeWindow(name, round(resize*W), round(resize*H))


def dicom_convert_jpg(dicom_path, jpg_path):
    if not os.path.exists(dicom_path):
        return -1

    ds = pydicom.dcmread(dicom_path)

    # Convert to float to avoid overflow or underflow losses.
    image_2d = ds.pixel_array.astype(float)

    # Rescaling grey scale between 0-255
    image_2d_scaled = (np.maximum(image_2d, 0) / image_2d.max()) * 255.0

    # Convert to uint
    image_2d_scaled = np.uint8(image_2d_scaled)
    #cv2.imwrite(jpg_path, image_2d_scaled)

    return 0


def create_hu_lookup(bits):
    lookup = [i for i in range(1<<bits)]
    return lookup


def read_dcm_by_pydicom(filename):
    dcm = pydicom.dcmread(filename)
    # try:
    #     dcm = pydicom.dcmread(fileanme)
    # except Exception as e:
    #     print('Error:', e, fileanme)
    #     return dcm_info

    try:
        dcm.decompress()
    except Exception as e:
        print('Error:', e, filename)
        return None
    return dcm

def read_dcm_info(filename):
    """
    read dicom file
    :param fileanme:dicom file name
    :param dcm_info:dict
    :return:dict
    """
    dcm_info = {}
    dcm_info['image'] = None

    if not os.path.exists(filename):
        return dcm_info

    dcm = read_dcm_by_pydicom(filename)
    if dcm is None:
        return dcm_info

    dcm_info['dcm_file'] = str(filename)
    dcm_info['Modality'] = dcm.Modality
    dcm_info['PatientSex'] = dcm.PatientSex

    if Tag(BODY_PART_TAG) in dcm:
        body_name = dcm.BodyPartExamined
        if body_name not in BODY_PART_NAME:
            encode_name = body_name.encode(encoding="iso8859")
            body_name = encode_name.decode('gbk')
        dcm_info['body_part'] = body_name # str(dcm[BODY_PART_TAG[0], BODY_PART_TAG[1]]._value).upper()
    else:
        dcm_info['body_part'] = ''

    # if Tag(PIXEL_SPACING_TAG) in dcm:
    #     dcm_info['spacing'] = list(map(float, dcm[PIXEL_SPACING_TAG[0], PIXEL_SPACING_TAG[1]]._value))
    # else:
    #     dcm_info['spacing'] = None

    if Tag(PHOTOMETRIC_INTERPRETATION_TAG) in dcm:
        dcm_info['photo_metric_interpretation'] = dcm[PHOTOMETRIC_INTERPRETATION_TAG[0], PHOTOMETRIC_INTERPRETATION_TAG[1]]._value
    else:
        dcm_info['photo_metric_interpretation'] = ''

    if Tag(BITS_STORED_TAG) in dcm:
        dcm_info['bits'] = int(dcm[BITS_STORED_TAG[0], BITS_STORED_TAG[1]]._value)
    else:
        dcm_info['bits'] = 8

    if Tag(RESCALE_SLOPE_TAG) in dcm:
        dcm_info['rescaleslope'] = int(dcm[RESCALE_SLOPE_TAG[0], RESCALE_SLOPE_TAG[1]]._value)
    else:
        dcm_info['rescaleslope'] = None

    if Tag(RESCALE_INTERCEPT_TAG) in dcm:
        dcm_info['rescaleintercept'] = int(dcm[RESCALE_INTERCEPT_TAG[0], RESCALE_INTERCEPT_TAG[1]]._value)
    else:
        dcm_info['rescaleintercept'] = None

    if hasattr(dcm, 'Rows'):
        dcm_info['rows'] = int(dcm.Rows)
    else:
        dcm_info['rows'] = 0
    if hasattr(dcm, 'Columns'):
        dcm_info['cols'] = int(dcm.Columns)
    else:
        dcm_info['cols'] = 0

    if hasattr(dcm, 'WindowCenter'):
        dcm_info['wc'] = int(dcm.WindowCenter)
    else:
        dcm_info['wc'] = 0
    if hasattr(dcm, 'WindowWidth'):
        dcm_info['ww'] = int(dcm.WindowWidth)
    else:
        dcm_info['ww'] = 0

    if hasattr(dcm, 'PatientOrientation'):
        dcm_info['PatientOrientation'] = str(dcm.PatientOrientation).upper()
    else:
        dcm_info['PatientOrientation'] = None

    if hasattr(dcm, 'pixel_array'):
        dcm_info['image'] = dcm.pixel_array
    else:
        dcm_info['image'] = None

    return dcm_info


def convert_dcm_grayscale8(dcm_info, norm=False):
    """
    convert ct value to gray 8bit image
    :param dcm_info:
    :return:
    """
    img_data = None
    if dcm_info['image'] is not None:
        img_data = dcm_info['image'].astype(float)
        if dcm_info['wc'] != 0 and dcm_info['ww'] != 0:
            img_data = convert_fixed_ww_wc(img_data, dcm_info['ww'], dcm_info['wc'])
        else:
            img_data = convert_defalut_ww_wc(img_data)

        if not norm:
            img_data = np.uint8(img_data*255)
        if 'photo_metric_interpretation' in dcm_info and dcm_info['photo_metric_interpretation'].upper().find('MONOCHROME1') >= 0:
            if not norm:
                img_data = (255 - img_data)
            else:
                img_data = 1.0 - img_data
    return img_data


def ajuste_dcm_grayscale8(dcm_info, ww, wc):
    img_data = dcm_info['image'].astype(float)
    img_type = dcm_info['Modality'].upper()

    if ww != 0 and wc != 0:
        img_data = convert_fixed_ww_wc(img_data, ww, wc)
    else:
        img_data = convert_defalut_ww_wc(img_data)
    img_data = np.uint8(img_data * 255)
    if dcm_info['photo_metric_interpretation'].upper().find('MONOCHROME1') >= 0:
        img_data = (255 - img_data)

    return img_data


def convert_defalut_ww_wc(img_data):
    img_temp = img_data
    img_temp = (np.maximum(img_temp, 0) / img_temp.max()) #* 255.0
    return img_temp#np.uint8(img_temp)


def convert_fixed_ww_wc(img_data, ww, wc):
    img_temp = img_data
    img_temp.flags.writeable = True
    min_v = (2 * wc - ww) / 2.0 + 0.5
    max_v = (2 * wc + ww) / 2.0 + 0.5
    factor = 1.0 / (max_v - min_v) # 255.0 / (max_v - min_v)

    img_temp = (img_temp - min_v) * factor

    min_index = img_temp < 0
    img_temp[min_index] = 0
    max_index = img_temp > 1#255
    img_temp[max_index] = 1#255
    return img_temp


def getSlicesPixelHu(slices):
    """
    convert 2d dicom slices into CT value(int16)
    :param slices:2d dicom slices
    :return:3D CT image(n*w*h), spacing
    """
    image = np.stack([s.pixel_array for s in slices])

    # Convert to int16 (from sometimes int16),
    # should be possible as values should always be low enough (<32k)
    image = image.astype(np.int16)

    # Convert to Hounsfield units (HU)
    for slice_number in range(len(slices)):
        intercept = slices[slice_number].RescaleIntercept
        slope = slices[slice_number].RescaleSlope

        if slope != 1:
            image[slice_number] = slope * image[slice_number].astype(np.float64)
            image[slice_number] = image[slice_number].astype(np.int16)

        image[slice_number] += np.int16(intercept)

    return np.array(image, dtype=np.int16), np.array([slices[0].SliceThickness] + slices[0].PixelSpacing._list,
                                                     dtype=np.float32)


def convert_pixel_to_ct(dcm_info):
    if 'rescaleslope' in dcm_info and  'rescaleintercept' in dcm_info and dcm_info['rescaleslope'] and dcm_info['rescaleintercept'] and dcm_info['image']:
        dcm_info['image'] = dcm_info['image'] * np.int16(dcm_info['rescaleslope']) + np.int16(dcm_info['rescaleintercept'])


def read_dcm_to_8bit(file):
    """
    Read dicom ct value convert grayscale 8bit image.
    :param file:
    :return:
    """
    dcm_info = read_dcm_info(file)
    convert_pixel_to_ct(dcm_info)
    return convert_dcm_grayscale8(dcm_info)


def read_dcm_to_norm(file):
    """
    Read dicom ct value convert grayscale normalized image.
    :param file:
    :return:
    """
    dcm_info = read_dcm_info(file)
    convert_pixel_to_ct(dcm_info)
    return convert_dcm_grayscale8(dcm_info, True)


def read_dcm_to_sitk(file):
    reader = sitk.ImageFileReader()
    reader.SetImageIO("GDCMImageIO")
    reader.SetFileName(file)
    image = reader.Execute()

    # image = sitk.ReadImage(name)
    wc_tag = '0028|1050'
    ww_tag = '0028|1051'
    body_tag = '0018|0015'
    wc = ww = 0
    # print(image.GetMetaDataKeys())
    if wc_tag in image.GetMetaDataKeys():
        wc = int(image.GetMetaData(wc_tag))
    if ww_tag in image.GetMetaDataKeys():
        ww = int(image.GetMetaData(ww_tag))
    if body_tag in image.GetMetaDataKeys():
        body = image.GetMetaData(body_tag)

    dcm_info = {}
    dcm_info['ww'] = ww
    dcm_info['wc'] = wc
    dcm_info['image'] = sitk.GetArrayFromImage(image)

    return convert_dcm_grayscale8(dcm_info)


def lumTrans(img):
    lungwin = np.array([-1200.,600.])
    newimg = (img-lungwin[0])/(lungwin[1]-lungwin[0])
    newimg[newimg<0]=0
    newimg[newimg>1]=1
    newimg = (newimg*255).astype('uint8')
    return newimg


def normalize_ct_to_unit(sitkImage, shift, scale):
    """
    normalize sitkImage
    :param sitkImage: sitkImage
    :return:normalized sitkImage
    """
    sitkImage = sitk.Cast(sitkImage, sitk.sitkFloat64)
    normalized_sitkImage = sitk.ShiftScale(sitkImage, shift=shift, scale=scale)
    img_array = sitk.GetArrayFromImage(normalized_sitkImage)
    img_array[img_array < 0.0] = 0.0
    img_array[img_array > 1.0] = 1.0
    origin = normalized_sitkImage.GetOrigin()
    spacing = normalized_sitkImage.GetSpacing()
    try:
        sitkImage = sitk.GetImageFromArray(img_array)
    except Exception as e:
        print(e)
        return None
    sitkImage.SetOrigin(origin)
    sitkImage.SetSpacing(spacing)

    return sitkImage


# normalize lung from [-1100 200] to [0 1]
# sitk.Normalize() is used
def normalize_ct_lung_to_unit(sitkImage):
    _shift = 1100.0
    _scale = 1.0 / 1300.0
    return normalize_ct_to_unit(sitkImage, _shift, _scale)


def load_mhd_image(filename):
    """
    simple itk load mhd file
    :param filename:inpute filename
    :return:dict
    """
    with open(filename) as f:
        contents = f.readlines()
        line = [k for k in contents if k.startswith('TransformMatrix')][0]
        transformM = np.array(line.split(' = ')[1].split(' ')).astype('float')
        transformM = np.round(transformM)
        if np.any(transformM != np.array([1, 0, 0, 0, 1, 0, 0, 0, 1])):
            isflip = True
        else:
            isflip = False
    itkimage = sitk.ReadImage(filename)
    numpyOrigin = np.array(list(reversed(itkimage.GetOrigin())))
    numpySpacing = np.array(list(reversed(itkimage.GetSpacing())))
    return {"itkimage": itkimage, "origin": numpyOrigin, "spacing": numpySpacing, "flip": isflip}


def get_numpy_from_sitk(sitkimage):
    """
    convert sitkimage to numpy array
    :param sitkimage:
    :return:
    """
    return sitk.GetArrayFromImage(sitkimage)


def fn_timer(function):
  @functools.wraps(function)
  def function_timer(*args, **kwargs):
    t0 = time.time()
    result = function(*args, **kwargs)
    t1 = time.time()
    print ("Total time running %s: %s seconds" %
        (getattr(function,'__name__'), str(t1-t0))
        )
    return result
  return function_timer

@fn_timer
def test_pydicom():
    name = '/home/blake/data/dataset/datasets/DR/data_pool/background/0520181004344235-0000000001-0002-10001-1.2.410.200048.8059.20181004082631.1.2.2.dcm'
    for i in range(100):
        dcm_info = read_dcm_info(name)
        print(dcm_info['body_part'])
        img = convert_dcm_grayscale8(dcm_info)
@fn_timer
def test_sitk():
    name = '/home/blake/data/dataset/datasets/DR/data_pool/background/0520181004344235-0000000001-0002-10001-1.2.410.200048.8059.20181004082631.1.2.2.dcm'
    for i in range(100):
        img_data = None
        reader = sitk.ImageFileReader()
        reader.SetImageIO("GDCMImageIO")
        reader.SetFileName(name)
        image = reader.Execute()

        #image = sitk.ReadImage(name)
        wc_tag = '0028|1050'
        ww_tag = '0028|1051'
        body_tag = '0018|0015'
        wc = ww = 0
        #print(image.GetMetaDataKeys())
        if wc_tag in image.GetMetaDataKeys():
            wc = int(image.GetMetaData(wc_tag))
        if ww_tag in image.GetMetaDataKeys():
            ww = int(image.GetMetaData(ww_tag))
        if body_tag in image.GetMetaDataKeys():
            body = image.GetMetaData(body_tag)
        image_array = sitk.GetArrayFromImage(image)

        if wc !=0 and ww != 0:
            img_data = convert_fixed_ww_wc(image_array[0], ww, wc)
        else:
            img_data = convert_defalut_ww_wc(image_array[0])

        # plt.imshow(img_data, cmap=plt.cm.gray)
        # plt.show()

def remove_privacy(dcm_file):
    convert_dcm = read_dcm_by_pydicom(dcm_file)
    convert_dcm.PatientName = ""
    convert_dcm.InstitutionName = ""
    convert_dcm.PatientBirthDate = ""

    save_path = 'convert_mm.dcm'
    convert_dcm.save_as(save_path)


def main():
    import matplotlib.pyplot as plt
    path = '/home/blake/medical/DR/QC/gdDR/data/77/63186359.dcm'
    #path = '/home/blake/data/medical/datasets/breast/CBIS-DDSM/Mass-Training_P_00384_RIGHT_MLO/07-20-2016-DDSM-02609/1-full mammogram images-55971/000000.dcm'
    dn = 'E:\\medical\\breast\\datasets\\乳腺癌_convert\\001091461200\\1241581-王小琴-20190424342445\\0000001_convert.dcm'
    #dn = 'E:\\medical\\breast\\datasets\\CBIS\\CBIS_Demo\\1241581-王小琴-20190424342445\\000002.dcm'

    dcm = read_dcm_info(dn)
    # cv_show('image', np.squeeze(dcm["image"]), 0.25)
    # cv2.waitKey(0)

if __name__ == '__main__':
    main()



